package com.weilus.dataflow;

import lombok.extern.slf4j.Slf4j;
import org.springframework.batch.core.Job;
import org.springframework.batch.core.Step;
import org.springframework.batch.core.configuration.annotation.EnableBatchProcessing;
import org.springframework.batch.core.configuration.annotation.JobBuilderFactory;
import org.springframework.batch.core.configuration.annotation.StepBuilderFactory;
import org.springframework.batch.core.configuration.annotation.StepScope;
import org.springframework.batch.core.launch.JobLauncher;
import org.springframework.batch.core.partition.PartitionHandler;
import org.springframework.batch.core.partition.support.TaskExecutorPartitionHandler;
import org.springframework.batch.item.ExecutionContext;
import org.springframework.batch.item.database.JdbcCursorItemReader;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.autoconfigure.jdbc.DataSourceProperties;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.boot.jdbc.DataSourceBuilder;
import org.springframework.cloud.task.configuration.EnableTask;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.task.SimpleAsyncTaskExecutor;
import org.springframework.jdbc.core.ArgumentPreparedStatementSetter;

import javax.sql.DataSource;
import java.util.HashMap;
import java.util.Map;

@Slf4j
@Configuration
@EnableTask
@EnableBatchProcessing
public class BatchConfiguration {
    @Autowired
    JobBuilderFactory jobBuilderFactory;
    @Autowired
    StepBuilderFactory stepBuilderFactory;
    @Autowired
    JobLauncher jobLauncher;

    @Bean("taskDataSourceProperties")
    @ConfigurationProperties(prefix = "task.datasource")
    public DataSourceProperties taskDataSourceProperties(){
        return new DataSourceProperties();
    }

    @Bean("taskDataSource")
    public DataSource dataSource(@Qualifier("taskDataSourceProperties") DataSourceProperties properties){
        return DataSourceBuilder.create()
                .driverClassName(properties.getDriverClassName())
                .username(properties.getUsername())
                .password(properties.getPassword())
                .url(properties.getUrl())
                .build();
    }

    @Bean
    @StepScope
    public JdbcCursorItemReader<User> reader(
            @Value("#{stepExecutionContext[startId]}") Integer startId,
            @Value("#{stepExecutionContext[endId]}")Integer endId,
            @Qualifier("taskDataSource") DataSource dataSource){
        log.debug("startId= {} endId= {}",startId,endId);
        JdbcCursorItemReader<User> reader = new JdbcCursorItemReader();
        reader.setSql("SELECT * FROM user WHERE id >=? AND id < ?");
        reader.setRowMapper((rs,i)->new User(rs.getInt("id"),rs.getString("name")));
        reader.setDataSource(dataSource);
        reader.setPreparedStatementSetter(new ArgumentPreparedStatementSetter(new Object[]{startId,endId}));
        return reader;
    }


    @Bean
    public Step slaveStep(JdbcCursorItemReader<User> reader){
        return this.stepBuilderFactory.get("slaveStep")
                .chunk(3)
                .reader(reader)
                .writer(user->log.debug(user.toString()))
                .allowStartIfComplete(true)
                .build();
    }

    @Bean
    public PartitionHandler partitionHandler(){
        TaskExecutorPartitionHandler handler = new TaskExecutorPartitionHandler();
        handler.setStep(slaveStep(reader(null,null,null)));
        handler.setGridSize(2);
        handler.setTaskExecutor(new SimpleAsyncTaskExecutor());
        return handler;
    }

    @Bean
    public Step masterStep(){
        return this.stepBuilderFactory.get("masterStep")
                .partitioner("slaveStep",gridSize->{
                    System.out.println("gridSize ............"+gridSize);
                    Map<String, ExecutionContext> map = new HashMap(gridSize);
                    ExecutionContext ctx1 = new ExecutionContext(),ctx2 = new ExecutionContext();
                    ctx1.putLong("startId",1);
                    ctx1.putLong("endId",3);
                    ctx2.putLong("startId",3);
                    ctx2.putLong("endId",5);
                    map.put("partition1", ctx1);
                    map.put("partition2", ctx2);
                    return map;
                })
                .partitionHandler(partitionHandler())
                .allowStartIfComplete(true)
                .build();
    }

    @Bean
    public Job job(){
        return this.jobBuilderFactory.get("endOfDay")
                .start(masterStep())
                .build();
    }
}
